import torch
import yaml
from torch import nn
from config.module import CBL, ResUnit, DownSample
from config import cfg


class DarkNet53(nn.Module):
    def __init__(self):
        super().__init__()
        self.input_layer = nn.Sequential(
            CBL(3, 32, 3, 1)
        )

        layers = []
        with open(cfg.DARKNET35_PARAM_PATH, 'r', encoding='utf-8') as file:
            dic = yaml.safe_load(file)
            channels = dic['channels']
            block_nums = dic['block_nums']

        for idx, block_num in enumerate(block_nums):
            layers.append(self.make_layer(channels[idx], channels[idx + 1], block_num))
        self.hidden_layer = nn.Sequential(*layers)

    def make_layer(self, c_in, c_out, block_num):
        units = [DownSample(c_in, c_out)]
        for _ in range(block_num):
            units.append(ResUnit(c_out))
        return nn.Sequential(*units)

    def forward(self, x):
        x = self.input_layer(x)
        unit52_out = self.hidden_layer[:3](x)
        unit26_out = self.hidden_layer[3](unit52_out)
        unit13_out = self.hidden_layer[4](unit26_out)
        return unit52_out, unit26_out, unit13_out


if __name__ == '__main__':
    data = torch.randn(1, 3, 416, 416)
    net = DarkNet53()
    # out = net(data)
    # print(out.shape)
    # # torch.Size([1, 1024, 13, 13])
    outs = net(data)
    for out in outs:
        print(out.shape)
    # torch.Size([1, 256, 52, 52])
    # torch.Size([1, 512, 26, 26])
    # torch.Size([1, 1024, 13, 13])

    # darknet_hidden_param = {
    #     'channels': [32, 64, 128, 256, 512, 1024],
    #     'block_nums': [1, 2, 8, 8, 4]
    # }
    # with open('data.yaml', 'r', encoding='utf-8') as file:
    #     # yaml.safe_dump(darknet_hidden_param, file)
    #     dic = yaml.safe_load(file)
    #     channels = dic['channels']
    #     block_nums = dic['block_nums']
    # print(dic)
    # # {'block_nums': [1, 2, 8, 8, 4], 'channels': [32, 64, 128, 256, 512, 1024]}
    # print(channels)
    # # [32, 64, 128, 256, 512, 1024]
    # print(block_nums)
    # # [1, 2, 8, 8, 4]
    pass
